import onnx
from onnx2keras import onnx_to_keras
import tensorflow as tf
import torch
import torchvision
import torchvision.transforms as transforms
import shutil
import numpy as np

onnx_model = onnx.load('data\\mnist_net.onnx')
k_model = onnx_to_keras(onnx_model=onnx_model, input_names=['input'], change_ordering=True)

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize(0.5, 0.5)])
testset = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    download=True,
    transform=transform)
testloader = torch.utils.data.DataLoader(
    testset,
    batch_size=1,
    shuffle=False,
    num_workers=0)

# print('TEST')
# for image, label in testloader:
#     image = np.ascontiguousarray(np.transpose(image.numpy(), (0, 2, 3, 1)))
#     y = k_model.predict(image)[0]
#     print('label={}, y={}'.format(label, y.argmax()))

shutil.rmtree('saved_model', ignore_errors=True)
tf.saved_model.save(k_model, 'saved_model')
